package edu.buaa.common.server;

import com.alibaba.fastjson.JSON;
import com.google.common.base.Preconditions;
import edu.buaa.common.RuntimeEnv;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.common.transaction.AbstractTransaction.Result;
import edu.buaa.common.transaction.CreateTGraphAggrMaxIndexTx;
import edu.buaa.common.transaction.CreateTGraphTemporalValueIndexTx;
import edu.buaa.common.transaction.*;
import edu.buaa.common.utils.PFieldList;
import edu.buaa.common.utils.PVal;
import edu.buaa.utils.Helper;
import edu.buaa.utils.Pair;
import edu.buaa.utils.Triple;
import org.act.temporalProperty.query.TimePointL;
import org.act.temporalProperty.query.aggr.AggregationIndexQueryResult;
import org.act.temporalProperty.query.aggr.AggregationQuery;
import org.act.temporalProperty.query.aggr.ValueGroupingMap;
import org.act.temporalProperty.util.Slice;
import org.neo4j.graphdb.*;
import org.neo4j.temporal.*;
import org.neo4j.tooling.GlobalGraphOperations;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import static org.act.temporalProperty.index.IndexType.AGGR_MAX;

public class TGraphKernelServer extends DBSocketServer.ReqExecutor {
    private static final Logger log = LoggerFactory.getLogger(TGraphKernelServer.class);

    public TGraphKernelServer(boolean useIndex) {
        this.useIndex = useIndex;
    }

    public static void main(String[] args){
        boolean index = !"false".equalsIgnoreCase(System.getenv("NO_INDEX"));
        DBSocketServer server = new DBSocketServer( dbDir(), new TGraphKernelServer(index), index?8438:8439);
        RuntimeEnv env = RuntimeEnv.getCurrentEnv();
        String serverCodeVersion = env.name() + "." + Helper.codeGitVersion();
        System.out.println("server code version: "+ serverCodeVersion);
        try {
            server.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static File dbDir(){
        String path = Helper.mustEnv("DB_PATH");
        Preconditions.checkNotNull(path, "need arg: DB_PATH");
        File dbDir = new File(path);
        if( !dbDir.exists()){
            if(dbDir.mkdirs()) return dbDir;
            else throw new IllegalArgumentException("invalid dbDir");
        }else if( !dbDir.isDirectory()){
            throw new IllegalArgumentException("invalid dbDir");
        }
        return dbDir;
    }

    public enum EdgeType implements RelationshipType {
        EDGE_TO
    }

    private final Map<String, Long> nodeIdMap = new HashMap<>();
    private final Map<String, Long> relIdMap = new HashMap<>();
    private TemporalIndexMetaData aggrMaxIndex = null;
    private TemporalIndexMetaData aggrDurIndex = null;
    private TemporalIndexMetaData tvCondIndex = null;
    private final boolean useIndex;

    private long nodeId(String rawId){
        Long r = nodeIdMap.get(rawId);
        if(r==null){
            throw new IllegalArgumentException("node id "+rawId+" not exist in db.");
        }
        return r;
    }

    private long relId(String rawId){
        Long r = relIdMap.get(rawId);
        if(r==null){
            throw new IllegalArgumentException("rel id "+rawId+" not exist in db.");
        }
        return r;
    }

    @Override
    public void setDB(GraphDatabaseService db){
        this.db = db;
        try(Transaction tx = db.beginTx()){
            for(Node node: GlobalGraphOperations.at(db).getAllNodes()){
                String rawId = (String) node.getProperty("id");
                Preconditions.checkNotNull(rawId,"should not happen: id is null");
                nodeIdMap.put(rawId, node.getId());
//                System.out.println(rawId+" , "+ node.getId());
            }
            for(Relationship rel: GlobalGraphOperations.at(db).getAllRelationships()){
                String rawId = (String) rel.getProperty("id");
                Preconditions.checkNotNull(rawId,"should not happen: id is null");
                relIdMap.put(rawId, rel.getId());
//                System.out.println(rawId+" , "+ node.getId());
            }
            for(TemporalIndexMetaData m : db.temporalIndex().allIndex()){
                System.out.println(m);
                if(m.isOnline()) {
                    if (m.getIndexType().equalsIgnoreCase("AGGR_DURATION")) aggrDurIndex = m;
                    if (m.getIndexType().equalsIgnoreCase("AGGR_MAX")) aggrMaxIndex = m;
                    if (m.getIndexType().equalsIgnoreCase("SINGLE_VALUE")) tvCondIndex = m;
                }
            }
            System.out.println("AGGR_DUR_INDEX_ID: "+ aggrDurIndex);
            System.out.println("AGGR_MAX_INDEX_ID: "+ aggrMaxIndex);
            System.out.println("TIME_VAL_INDEX: "+tvCondIndex);
            tx.success();
        }
    }

    @Override
    public Result execute(String line) throws RuntimeException {
        AbstractTransaction tx = JSON.parseObject(line, AbstractTransaction.class);
        try {
            switch (tx.getTxType()) {
                case tx_import_static_data:
                    return execute((ImportStaticDataTx) tx);
                case tx_import_temporal_data:
                    return execute((ImportTemporalDataTx) tx);
                case tx_update_temporal_data:
                    return execute((UpdateTemporalDataTx) tx);
                case tx_query_snapshot:
                    return execute((SnapshotQueryTx) tx);
                case tx_query_entity_history:
//                    System.out.println(line);
                    return execute((EntityHistoryTx) tx);
                case tx_query_road_by_temporal_condition:
                    EntityTemporalConditionTx etx = (EntityTemporalConditionTx) tx;
                    if (useIndex && tvCondIndex != null && tvCondIndex.isNode() == etx.isNode() &&
                            etx.getP().equals(tvCondIndex.getProps().get(0).getLeft()) &&
                            tvCondIndex.getTimeStart().valInt() <= etx.getT0() && etx.getT1() <= tvCondIndex.getTimeEnd().valInt()) {
                        return execute(etx, 0);
                    } else {
                        return execute(etx);
                    }
                case tx_query_snapshot_aggr_max:
                    SnapshotAggrMaxTx stx = (SnapshotAggrMaxTx) tx;
                    if (useIndex && aggrMaxIndex != null && aggrMaxIndex.isNode() == stx.isNode() &&
                            stx.getP().equals(aggrMaxIndex.getProps().get(0).getLeft()) &&
                            aggrMaxIndex.getTimeStart().valInt() <= stx.getT0() && stx.getT1() <= aggrMaxIndex.getTimeEnd().valInt()) {
                        return execute(stx, aggrMaxIndex.getId());
                    } else return execute(stx);
                case tx_query_snapshot_aggr_duration:
                    SnapshotAggrDurationTx dtx = (SnapshotAggrDurationTx) tx;
                    if (useIndex && aggrDurIndex != null && aggrDurIndex.isNode() == dtx.isNode() &&
                            dtx.getP().equals(aggrDurIndex.getProps().get(0).getLeft()) &&
                            aggrDurIndex.getTimeStart().valInt() <= dtx.getT0() && dtx.getT1() <= aggrDurIndex.getTimeEnd().valInt()){
                        return execute(dtx, aggrDurIndex.getId());
                    }else return execute(dtx);
                case tx_index_tgraph_aggr_max:
                    return execute((CreateTGraphAggrMaxIndexTx) tx);
                case tx_index_tgraph_aggr_duration:
                    return execute((CreateTGraphAggrDurationIndexTx) tx);
                case tx_index_tgraph_temporal_condition:
                    return execute((CreateTGraphTemporalValueIndexTx) tx);
                default:
                    throw new UnsupportedOperationException();
            }
        } catch (Exception e) {
            int end = Math.min(60, line.length());
            log.error("ERROR processing TX: "+line.substring(0, end), e);
            throw new TransactionFailedException(e);
        }
    }

    private Result execute(ImportStaticDataTx tx){
        try(Transaction t = db.beginTx()) {
            PFieldList nodesData = tx.getNodes();
            Set<String> props = nodesData.keysWithout("id");
            int nSize = nodesData.size();
            for (int i=0; i<nSize; i++) {
                Node n = db.createNode();
                String id = nodesData.get("id", i).s();
                n.setProperty("id", id);
                for(String key : props){
                    n.setProperty(key, nodesData.get(key, i).getVal());
                }
                nodeIdMap.put(id, n.getId());
            }
            PFieldList relData = tx.getRels();
            props = relData.keysWithout("from", "to");
            int rSize = relData.size();
            for (int i=0; i<rSize; i++) {
                String fromId = relData.get("from", i).s();
                String toId = relData.get("to", i).s();
                Node start = db.getNodeById(nodeId(fromId));
                Node end = db.getNodeById(nodeId(toId));
                Relationship r = start.createRelationshipTo(end, EdgeType.EDGE_TO);
                for(String key : props){
                    r.setProperty(key, relData.get(key, i).getVal());
                }
                relIdMap.put(relData.get("id", i).s(), r.getId());
            }
            t.success();
        }
        return new Result();
    }

    private Result execute(ImportTemporalDataTx tx) {
        try(Transaction t = db.beginTx()) {
            PFieldList data = tx.getData();
            Set<String> props = data.keysWithout("id", "time");
            int tSize = data.size();
            for (int i=0; i<tSize; i++) {
                String id = data.get("id", i).s();
                TimePoint time = Helper.time(data.get("time", i).i());
                for(String prop : props){
                    PropertyContainer p = tx.isNode() ? db.getNodeById(nodeIdMap.get(id)) : db.getRelationshipById(relIdMap.get(id));
                    p.setTemporalProperty(prop, time, data.get(prop, i));
                }
            }
            t.success();
        }
        return new Result();
    }

    private Result execute(UpdateTemporalDataTx tx) {
        try(Transaction t = db.beginTx()) {
            PFieldList data = tx.getData();
            Set<String> props = data.keysWithout("id", "start", "end");
            int tSize = data.size();
            for (int i=0; i<tSize; i++) {
                String id = data.get("id", i).s();
                TimePoint s = Helper.time(data.get("start", i).i());
                TimePoint e = Helper.time(data.get("end", i).i());
                for(String prop : props){
                    PropertyContainer p = tx.isNode() ? db.getNodeById(nodeIdMap.get(id)) : db.getRelationshipById(relIdMap.get(id));
                    p.setTemporalProperty(prop, s, e, data.get(prop, i));
                }
            }
            t.success();
        }
        return new Result();
    }

    private Result execute(EntityHistoryTx tx) {
        TimePoint begin = Helper.time(tx.getBeginTime()), end = Helper.time(tx.getEndTime());
        try (Transaction t = db.beginTx()) {
            List<Triple<Integer, Integer, PVal>> answers = new ArrayList<>();
            PropertyContainer p;
            if(tx.isNode()){
                p = db.getNodeById(nodeId(tx.getEntity()));
            }else{
                p = db.getRelationshipById(relIdMap.get(tx.getEntity()));
            }
            p.getTemporalProperty(tx.getProp(), begin, end, new TimeIntervalRangeQuery(begin, end){
                @Override
                public void onEntry(TimePointL beginTime, TimePointL endTime, Object val) {
                    if(val instanceof Float) {
                        answers.add(Triple.of(beginTime.valInt(), endTime.valInt(), PVal.v(val)));
                    }
                }
            });
            EntityHistoryTx.Result result = new EntityHistoryTx.Result();
            result.setHistory(answers);
            return result;
        }
    }

    @FunctionalInterface
    private interface PropertyContainerCallback {
        void run(PropertyContainer entity);
    }

    private void all(boolean node, PropertyContainerCallback body){
        if(node) {
            for(Node n : GlobalGraphOperations.at(db).getAllNodes()){
                body.run(n);
            }
        }else{
            for(Relationship r : GlobalGraphOperations.at(db).getAllRelationships()){
                body.run(r);
            }
        }
    }

    private Result execute(SnapshotQueryTx tx){
        try(Transaction t = db.beginTx()) {
            List<Pair<String, PVal>> answers = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                Object v = entity.getTemporalProperty(tx.getPropertyName(), Helper.time(tx.getTimestamp()));
                if(v!=null){
                    answers.add(Pair.of(id, PVal.v(v)));
                }
            });
            SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
            result.setStatus(answers);
            return result;
        }
    }

    private Result execute(SnapshotAggrMaxTx tx){
        try(Transaction t = db.beginTx()){
            List<Pair<String, PVal>> answers = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                Object v = entity.getTemporalProperty(tx.getP(), Helper.time(tx.getT0()), Helper.time(tx.getT1()), new TemporalRangeQuery.MaxFloat());
                if(v!=null){
                    answers.add(Pair.of(id, PVal.v(v))); //放入answers中
                }
            });
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(answers);
            return result;
        }
    }

    private Result execute(SnapshotAggrDurationTx tx) {
        TimePoint begin = Helper.time(tx.getT0()), end = Helper.time(tx.getT1());
        try (Transaction t = db.beginTx()) {
            List<Triple<String, PVal, Integer>> answers = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                Object v = entity.getTemporalProperty(tx.getP(), begin, end,
                        new TimeIntervalRangeQuery.Duration<PVal>(begin, end, tx.getIntStartTreeSet(), PVal::v));
                if(v instanceof HashMap){
                    HashMap<PVal, Integer> m = (HashMap<PVal, Integer>) v;
                    m.forEach((k, val) -> answers.add(Triple.of(id, k, val)));
                }
            });
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(answers);
            return result;
        }
    }

    private Result execute(EntityTemporalConditionTx tx) {
        try (Transaction t = db.beginTx()) {
            List<String> res = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                entity.getTemporalProperty(tx.getP(), Helper.time(tx.getT0()), Helper.time(tx.getT1()), new TemporalRangeQuery() {
                    @Override
                    public boolean onNewEntry(long entityId, int propertyId, TimePointL time, Object val) {
                        if (val instanceof Comparable && PVal.within(tx.getVMin(), true, PVal.v(val), tx.getVMax(), true)) {
                            res.add(id);
                            return false;
                        }
                        return true;
                    }
                });
            });
            EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
            result.setEntities(res);
            return result;
        }
    }

//    ====================== index: create and query =======================

    private Result execute(CreateTGraphTemporalValueIndexTx tx)  {
        long indexId = -1;
        try(Transaction t = db.beginTx()){
            indexId = db.temporalIndex().nodeCreateValueIndex(Helper.time(tx.getStart()), Helper.time(tx.getEnd()), tx.getProps().toArray(new String[]{}));
            t.success();
        }
        while(true) try {
            db.temporalIndex().awaitIndexOnline(indexId);
            break;
        } catch (InterruptedException e) {
            Thread.interrupted();
        }
        CreateTGraphTemporalValueIndexTx.Result r = new CreateTGraphTemporalValueIndexTx.Result();
        r.setIndexId(indexId);
        return r;
    }

    private Result execute(CreateTGraphAggrDurationIndexTx tx) {
        long indexId = -1;
        try(Transaction t = db.beginTx()){
            indexId = db.temporalIndex().nodeCreateDurationIndex(
                    Helper.time(tx.getStart()), Helper.time(tx.getEnd()),
                    tx.getProName(), tx.getEvery(), tx.getTimeUnit(), new ValueGroupingMap.FloatRangeGroupMap(tx.getIntervalStart()));
            t.success();
        }
        while(true) try {
            db.temporalIndex().awaitIndexOnline(indexId);
            break;
        } catch (InterruptedException e) {
            Thread.interrupted();
        }
        CreateTGraphAggrDurationIndexTx.Result r = new CreateTGraphAggrDurationIndexTx.Result();
        r.setIndexId(indexId);
        return r;
    }

    private Result execute(CreateTGraphAggrMaxIndexTx tx) {
        long indexId = -1;
        try(Transaction t = db.beginTx()){
            indexId = db.temporalIndex().nodeCreateMinMaxIndex(
                    Helper.time(tx.getStart()), Helper.time(tx.getEnd()),
                    tx.getProName(), tx.getEvery(), tx.getTimeUnit(), AGGR_MAX);
            t.success();
        }
        while(true) try {
            db.temporalIndex().awaitIndexOnline(indexId);
            break;
        } catch (InterruptedException e) {
            Thread.interrupted();
        }
        CreateTGraphAggrMaxIndexTx.Result r = new CreateTGraphAggrMaxIndexTx.Result();
        r.setIndexId(indexId);
        return r;
    }

    private Result execute(SnapshotAggrMaxTx tx, long indexId){
//        System.out.println("index query(" + tx.getT0() +"~"+tx.getT1()+") via max index");
        try(Transaction t = db.beginTx()){
            List<Pair<String, PVal>> answers = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                AggregationIndexQueryResult v = entity.getTemporalPropertyWithIndex(tx.getP(), Helper.time(tx.getT0()), Helper.time(tx.getT1()), indexId);
                if(v!=null){
                    Map<Integer, Slice> result = v.getMinMaxResult();
                    answers.add(Pair.of(id, result.get(AggregationQuery.MAX).input().readFloat())); // 0 min, 1 max
                }
            });
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(answers);
            return result;
        }
    }

    private Result execute(SnapshotAggrDurationTx tx, long indexId) {
//        System.out.println("index query(" + tx.getT0() +"~"+tx.getT1()+") via dur index");
        try (Transaction t = db.beginTx()) {
            List<Triple<String, PVal, Integer>> answers = new ArrayList<>();
            all(tx.isNode(), entity -> {
                String id = (String) entity.getProperty("id");
                AggregationIndexQueryResult v = entity.getTemporalPropertyWithIndex(tx.getP(), Helper.time(tx.getT0()), Helper.time(tx.getT1()), indexId);
                if(v!=null){
                    Map<Float, Integer> result = v.getDurationWithMapping();
                    result.forEach((k, val)-> answers.add(Triple.of(id, k, val)));
                }
            });
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(answers);
            return result;
        }
    }

    private Result execute(EntityTemporalConditionTx tx, long indexId){
        System.out.println("TGraphKernelServer: query(" + tx.getT0() +"~"+tx.getT1()+" "+ "via index("+tvCondIndex+")");
        List<String> answersFinal;
        try(Transaction t = db.beginTx()){
            TemporalIndexManager.PropertyValueIntervalBuilder query;
            if(tx.isNode()){
                query = db.temporalIndex().nodeQueryValueIndex(Helper.time(tx.getT0()), Helper.time(tx.getT1()));
            }else{
                query = db.temporalIndex().relQueryValueIndex(Helper.time(tx.getT0()), Helper.time(tx.getT1()));
            }
            query.propertyValRange(tx.getP(), tx.getVMin(), tx.getVMax());
            List<IntervalEntry> answers = query.query();
            answersFinal = answers.stream().map(IntervalEntry::getEntityId).distinct().map(id->{
                Node node = db.getNodeById(id);
                return (String) node.getProperty("id");
            }).collect(Collectors.toList());
        }
        EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
        result.setEntities(answersFinal);
        return result;
    }
}
